# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import functools
import inspect
import logging


class EventClock(object):
  def __init__(self):
    self._timestamp = None
    self._num_updates = 0
    self._callbacks = []

  def is_ready(self):
    return self._timestamp is not None

  @property
  def timestamp(self):
    if self.is_ready():
      return self._timestamp
    else:
      raise RuntimeError('Event clock not ready yet!')

  def now(self):
    if self.is_ready():
      return datetime.datetime.fromtimestamp(self._timestamp / 1e9)
    else:
      raise RuntimeError('Event clock not ready yet!')

  def update(self, timestamp: int):
    self._timestamp = timestamp
    self._num_updates += 1

    for callback in self._callbacks:
      if not self.is_disabled(callback):
        callback(timestamp)
        self.update_callback_info(callback)

  @staticmethod
  def is_disabled(callback):
    return callback.__info__['disabled']

  @staticmethod
  def update_callback_info(callback):
    info = callback.__info__
    info['num_calls'] += 1
    if info['call_once'] is True:
      info['disabled'] = True

  def register(self, callback, call_once=False):
    callback = self._check_and_cook_callback(callback)
    callback.__info__ = {
        'num_calls': 0,
        'call_once': call_once,
        'disabled': False,
    }
    self._callbacks.append(callback)

  @staticmethod
  def _check_and_cook_callback(callback):
    signature = inspect.signature(callback)
    if len(signature.parameters) == 1:

      def new_callback(timestamp: int):
        return callback(timestamp)

      return new_callback
    elif len(signature.parameters) == 0:
      # noinspection PyUnusedLocal
      def new_callback(timestamp: int):
        return callback()

      return new_callback
    else:
      raise RuntimeError('Invalid callback signature!')


event_clock = EventClock()


class EventClockScheduler(object):
  def __init__(self):
    self._logger = logging.getLogger(__name__)
    self.jobs = []
    event_clock.register(self.on_time_change)

  def every(self, period):
    return Job(period, self)

  def on_time_change(self):
    runnable_jobs = (job for job in self.jobs if job.should_run)
    for job in sorted(runnable_jobs):
      job.run()

  def add_periodic_task(
      self,
      period,
      job_func,
      *,
      immediate_first_run=True,
      **kwargs,
  ):
    def on_event_clock_ready():
      if immediate_first_run:
        job_func(**kwargs)
      self.every(period).do(job_func, **kwargs)

    event_clock.register(on_event_clock_ready, call_once=True)

  def clear(self):
    self.jobs.clear()


class Job(object):
  def __init__(self, period, scheduler=None):
    self.period = period
    self.job_func = None  # the job job_func to run
    self.last_run = None  # datetime of the last run
    self.next_run = None  # datetime of the next run
    self.scheduler = scheduler  # scheduler to register with

  def __lt__(self, other):
    return self.next_run < other.next_run

  def do(self, job_func, *args, **kwargs):
    self.job_func = functools.partial(job_func, *args, **kwargs)
    try:
      functools.update_wrapper(self.job_func, job_func)
    except AttributeError:
      pass
    self._schedule_next_run()
    self.scheduler.jobs.append(self)
    return self

  @property
  def should_run(self):
    return event_clock.timestamp >= self.next_run

  def run(self):
    ret = self.job_func()
    self.last_run = event_clock.timestamp
    self._schedule_next_run()
    return ret

  def _schedule_next_run(self):
    self.next_run = event_clock.timestamp + int(self.period * 1e9)


event_clock_scheduler = EventClockScheduler()
